from fastapi import APIRouter, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import os
import tempfile
from pdfminer.high_level import extract_text
from typing import Optional
import shutil

router = APIRouter(prefix="/api/pdfminer/pdf", tags=["PDF处理"])

# 确保输入和输出目录存在
INPUT_DIR = "input_files"
OUTPUT_DIR = "output_files"
os.makedirs(INPUT_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

@router.post("/extract")
async def extract_pdf_text(file: UploadFile = File(...), save_output: bool = False):
    """
    从上传的PDF文件中提取文本
    - file: 上传的PDF文件
    - save_output: 是否将结果保存到文件中
    """
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(status_code=400, detail="只支持PDF文件")
    
    # 保存上传的文件到临时文件
    temp_dir = tempfile.mkdtemp()
    temp_file_path = os.path.join(temp_dir, file.filename)
    
    try:
        with open(temp_file_path, "wb") as temp_file:
            content = await file.read()
            temp_file.write(content)
        
        # 保存到输入目录
        input_file_path = os.path.join(INPUT_DIR, file.filename)
        shutil.copy2(temp_file_path, input_file_path)
        
        # 提取文本
        text = extract_text(temp_file_path)
        
        # 如果需要保存到文件
        if save_output:
            output_filename = os.path.splitext(file.filename)[0] + ".txt"
            output_path = os.path.join(OUTPUT_DIR, output_filename)
            with open(output_path, "w", encoding="utf-8") as out_file:
                out_file.write(text)
            return JSONResponse(content={
                "message": "文本提取成功",
                "filename": file.filename,
                "output_file": output_filename,
                "text_preview": text[:500] + "..." if len(text) > 500 else text
            })
        
        return JSONResponse(content={
            "message": "文本提取成功",
            "filename": file.filename,
            "text": text
        })
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理PDF时出错: {str(e)}")
    
    finally:
        # 清理临时文件
        shutil.rmtree(temp_dir, ignore_errors=True)

@router.get("/files")
async def list_pdf_files():
    """列出输入目录中的所有PDF文件"""
    pdf_files = [f for f in os.listdir(INPUT_DIR) if f.lower().endswith('.pdf')]
    return {"pdf_files": pdf_files}

@router.get("/outputs")
async def list_output_files():
    """列出输出目录中的所有文本文件"""
    txt_files = [f for f in os.listdir(OUTPUT_DIR) if f.lower().endswith('.txt')]
    return {"output_files": txt_files} 